import numpy as np
from MyFrame.utils import is_iterable
from MyFrame.config import Config
import weakref
import itertools   # 放在文件头部
import heapq


class Function:
    _tie_breaker = itertools.count()     # 全局自增计数器
    def __call__(self, *inputs):
        inputs = [Function.as_variable(x) for x in inputs]     #将输入变量转换为Variable类型
        if Config.enable_backprop:     #如果启用反向传播,则保存输入变量和函数代数、id
            self.inputs = inputs     #元组形式
            self.generation = max([x.generation for x in inputs])     #设置函数的代数：最大输入代数
            self.tid = next(Function._tie_breaker)     #设置函数的id, 用于在同一代中确定函数的优先级
        xs = [x.data for x in inputs]     #提取输入数据
        ys = self.forward(*xs)
        if not is_iterable(ys):
            ys = (ys,)
        outputs = [Variable(Function.as_array(y)) for y in ys]
        if Config.enable_backprop:     #如果启用反向传播,则设置输出变量的连接关系
            for output in outputs:
                output.set_creator(self)
            self.outputs = [weakref.ref(output) for output in outputs]   #函数对每个output元素使用弱引用，并保存在列表中
        return outputs if len(outputs) > 1 else outputs[0]
    @staticmethod
    def as_array(input):
        if np.isscalar(input):
            output = np.array(input)
        else:
            output = input                
        return output
    @staticmethod
    def as_variable(input):
        if isinstance(input, Variable):
            return input
        else:
            return Variable(input)
    def forward(self, *xs):
        raise NotImplementedError()
    def backward(self, *grad_in_ys):
        raise NotImplementedError()

class Add(Function):
    def forward(self, x0, x1):
        return x0 + x1
    def backward(self, grad_output):
        return grad_output, grad_output
def add(x0, x1):     #x0为Variable类型，x1为数值或numpy数组
    x1 = Function.as_array(x1)
    return Add()(x0, x1)
def radd(x0, x1):
    x1 = Function.as_array(x1)
    return Add()(x1, x0)

class Mul(Function):
    def forward(self, x0, x1):
        return x0 * x1
    def backward(self, grad_output):
        return grad_output * self.inputs[1].data, grad_output * self.inputs[0].data
def mul(x0, x1):
    x1 = Function.as_array(x1)
    return Mul()(x0, x1)
def rmul(x0, x1):
    x1 = Function.as_array(x1)
    return Mul()(x1, x0)

class Neg(Function):
    def forward(self, x):
        return -x
    def backward(self, grad_output):
        return -grad_output
def neg(x):
    return Neg()(x)

class Sub(Function):
    def forward(self, x0, x1):
        return x0 - x1
    def backward(self, grad_output):
        return grad_output, -grad_output
def sub(x0, x1):
    x1 = Function.as_array(x1)
    return Sub()(x0, x1)
class RSub(Function):
    def forward(self, x0, x1):
        return x1 - x0     #x - self
    def backward(self, grad_output):
        return -grad_output, grad_output
def rsub(x0, x1):     #self, x
    x1 = Function.as_array(x1)
    return RSub()(x0, x1)     #上一次的bug找到了，这里写反了

class Div(Function):
    def forward(self, x0, x1):
        return x0 / x1
    def backward(self, grad_output):
        x0, x1 = self.inputs[0].data, self.inputs[1].data
        return grad_output / x1, -grad_output * x0 / (x1 ** 2)
def div(x0, x1):
    x1 = Function.as_array(x1)
    return Div()(x0, x1)
class RDiv(Function):
    def forward(self, x0, x1):
        return x1 / x0     #self / x
    def backward(self, grad_output):
        x0, x1 = self.inputs[0].data, self.inputs[1].data
        return -grad_output * x1 / (x0 ** 2), grad_output / x0
def rdiv(x0, x1):     #x, self
    x1 = Function.as_array(x1)
    return RDiv()(x0, x1)     #上一次的bug找到了，这里写反了

class Pow(Function):
    def forward(self, x0, x1):
        return x0 ** x1
    def backward(self, grad_output, eps=1e-10):
        x0, x1 = self.inputs[0].data, self.inputs[1].data
        return grad_output * x1 * x0 ** (x1 - 1), None
def pow(x0, x1):
    x1 = Function.as_array(x1)
    return Pow()(x0, x1)

class Variable: 
    def __init__(self, data, name=None):
        self.data = Function.as_array(data)
        self.grad = None
        self.creator = None
        self.generation = 0
        self.name = name
    def set_creator(self, func):
        self.creator = func
        self.generation = func.generation + 1     #设置输出变量代数：函数的代数+1
    def backward(self):
        funcs = []
        seen_set = set()     #记录已经计算过的函数,避免重复计算
        if self.grad is None:
            self.grad = np.ones_like(self.data)
        if self.creator is not None:
            seen_set.add(self.creator)
            heapq.heappush(funcs, (-self.creator.generation, self.creator.tid, self.creator))
        while funcs:
            _, _, func = heapq.heappop(funcs)
            for output in func.outputs:
                if output().grad is None:
                    output().grad = np.ones_like(output().data)
            out_grads = [output().grad for output in func.outputs]
            in_grads = func.backward(*out_grads)
            if not is_iterable(in_grads):     #变为可迭代对象
                in_grads = (in_grads,)
            for input, in_grad in zip(func.inputs, in_grads):
                if input.grad is None:
                    input.grad = in_grad
                else:
                    input.grad = input.grad + in_grad
                if input.creator is not None:
                    if input.creator not in seen_set:
                        seen_set.add(input.creator)
                        heapq.heappush(funcs, (-input.creator.generation, input.creator.tid, input.creator))
    def zero_grad(self):
        self.grad = None
    #获取Variable的属性
    @property
    def shape(self):
        return self.data.shape
    @property
    def ndim(self):
        return self.data.ndim
    @property
    def size(self):
        return self.data.size
    @property
    def dtype(self):
        return self.data.dtype
    @property
    def T(self):
        return Variable(self.data.T)
    def __len__(self):
        return len(self.data)
    def __str__(self):
        return "Variable(data={})".format(self.data)
    def __repr__(self):
        if self.data is None:
            return "Variable(None)"
        else:
            p = str(self.data).replace("\n", "\n" + " " * 9)
            return "Variable(" + p + ")"
        
